#include "bxdf.h"
#include "sampling.h"
#include "rng.h"

#include <time.h>

// BxDF utility function

//////////////////////////////////////////////
//	Medium			index of reflection
//	Vacuum			1.0
//	Air				1.00029
//	Ice				1.31
//	Water			1.333
//	Fused quartz	1.46
//	Glass			1.5-1.6
//	Sapphire		1.77
//	Diamond			2.42
///////////////////////////////////////////////

// Fr for dielectric, Ft = 1 - Fr
Spectrum FrDiel(float cosi, float cost, const Spectrum &etai, const Spectrum &etat) {
	Spectrum Rparl = (etat*cosi - etai*cost) / (etat*cosi + etai*cost);
	Spectrum Rperp = (etai*cosi - etat*cost) / (etai*cosi + etat*cost);
	return 0.5f * (Rparl * Rparl + Rperp * Rperp);
}

///////////////////////////////////////////////
//	Objects			eta			k
//	Gold			0.370		2.82
//	Silver			0.177		3.638
//	Copper			0.617		2.63
//	Steel			2.485		3.433
///////////////////////////////////////////////
// Fr for conductor, no Ft
Spectrum FrCond(float cosi, const Spectrum &eta, const Spectrum &k) {
	Spectrum tmp = (eta*eta + k*k) * cosi * cosi;
	Spectrum Rparl2 = (tmp - 2*eta*cosi + 1) / (tmp + 2*eta*cosi + 1);
	tmp = eta * eta + k * k;
	Spectrum Rperp2 = (tmp - 2*eta*cosi+cosi*cosi) / (tmp + 2*eta*cosi+cosi*cosi);
	return (Rparl2 + Rperp2) * 0.5f;
}
// Assume k = 0;
Spectrum FrApproxEta(const Spectrum &Fr) {
	Spectrum reflectance = Fr.Clamp(0.f,0.999f);
	return (Spectrum(1.0f) + reflectance.Sqrt()) / (Spectrum(1.0f) - reflectance.Sqrt());
}
// Assume eta = 1
Spectrum FrApproxiK(const Spectrum &Fr) {
	Spectrum reflectance = Fr.Clamp(0.f,0.999f);
	return 2.f * (reflectance / Spectrum(1.0f) - reflectance).Sqrt();
}

Spectrum FresnelDielectric::Evaluate(float cosi) const {
	cosi = Clamp(cosi, -1.f, 1.f);
	bool entering = cosi > 0;
	float ei = eta_i, et = eta_t;
	if(!entering)
		swap(ei, et);
	// Snell's law
	float sint = ei/et * sqrtf(max(0.f, 1.f - cosi*cosi));
	if(sint > 1.) {
		// Total internal reflection
		return 1.f;
	}
	else {
		float cost = sqrtf(max(0.f, 1.f-sint*sint));
		return FrDiel(fabsf(cosi), cost, ei, et);
	}
}
Spectrum BxDF::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf)  const {
	// Default to Lambertian(ideal diffuse)
	// Use consine-weighted sampling
	*wi = CosineSamplingHemisphere(BxDFSample);
	if (wo.z < 0.) wi->z *= -1.f;
	*pdf = Pdf(wo, *wi);
	return f(wo, *wi);
}

Spectrum BxDF::rho(int nSamples /* = 16 */, Point2 *samples /* = NULL */) const {
	RNG rng(time(NULL));
	bool needClear = false;
	if(samples == NULL) {
		// Create 16x2 samples for estimate rho
		// samples[2*i] for sampling wo
		// samples[2*i+1] for sampling wi
		samples = new Point2[2*nSamples];
		for(int i = 0;i < 2*nSamples; i++) {
			samples[i].x = rng.RandomFloat();
			samples[i].y = rng.RandomFloat();
		}
	}
	// Estimate rho with MC
	Spectrum r = 0.f, f;
	float pdf_o, pdf_i;
	Vector3 wo, wi;
	for(int i = 0; i < nSamples; i++) {
		pdf_o = 0.f; 
		pdf_i = 0.f;
		wo = UniformSamplingHemisphere(samples[2*i]);
		pdf_o = INV_TWOPI;
		f = Sample_f(wo, &wi, samples[2*i+1], &pdf_i);
		if(pdf_i > 0.f) {
			r += f * Frame::AbsCosTheta(wo) * Frame::AbsCosTheta(wi) / (pdf_o * pdf_i);
		}
	}

	if(needClear)
		SAFE_DELETE_ARRAY(samples);
	return r / (M_PI * nSamples);
}

Spectrum BxDF::rho(const Vector3 &wo, int nSamples /* = 16 */, Point2 *samples /* = NULL */) const {
	RNG rng(time(NULL));
	bool needClear = false;
	if(samples == NULL) {
		// Create 16 samples for estimate rho
		samples = new Point2[nSamples];
		for(int i = 0; i < nSamples; i++) {
			samples[i].x = rng.RandomFloat();
			samples[i].y = rng.RandomFloat();
		}
		needClear = true;
	}
	// estimate rho with MC
	Spectrum r = 0.f, f;
	float pdf;
	Vector3 wi;
	for(int i = 0; i < nSamples; i++) {
		pdf = 0.f;
		f = Sample_f(wo, &wi, samples[i], &pdf);
		if(pdf > 0.)
			r += f * Frame::AbsCosTheta(wi) / pdf;
	}

	if(needClear)
		SAFE_DELETE_ARRAY(samples);
	
	return r / (float)nSamples;
}

float BxDF::Pdf(const Vector3 &wo, const Vector3 &wi) const {
	// Default to Lambertian(ideal diffuse)
	if(Frame::SameHemisphere(wi, wo))
		return Frame::AbsCosTheta(wi) * INV_PI;
	return 0.f;
}

Spectrum BRDFToBTDF::f(const Vector3 &wo, const Vector3 &wi) const {
	return brdf->f(wo, Frame::OtherHemisphere(wi));
}

Spectrum BRDFToBTDF::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const {
	Spectrum f = brdf->Sample_f(wo, wi, BxDFSample, pdf);
	*wi = Frame::OtherHemisphere(*wi);
	return f;
}

Spectrum BRDFToBTDF::rho(int nSamples , Point2 *samples ) const {
	return brdf->rho(nSamples, samples);
}

Spectrum BRDFToBTDF::rho(const Vector3 &wo, int nSamples , Point2 *samples ) const {
	return brdf->rho(Frame::OtherHemisphere(wo), nSamples, samples);
}

Spectrum SpecularReflection::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const {
	*pdf = 1.f;
	*wi = Vector3(-wo.x, -wo.y, wo.z);
	// f = fr(wo) * R / abs(cos(wi))
	return fresnel->Evaluate(Frame::CosTheta(wo)) * R / Frame::AbsCosTheta(*wi);
}

Spectrum SpecularTransmission::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const {
	bool entering  = Frame::CosTheta(wo) > 0;
	float ei = eta_i, et = eta_t;
	if(!entering)
		swap(ei, et);
	// Compute transmit direction
	float sini2 = Frame::SinTheta2(wo);
	float eta = et / et;
	float sint2 = eta * eta * sini2;
	if(sint2 >= 1.)
		return 0.f;
	float cost = sqrtf(max(0.f, 1-sint2));
	if(entering) cost = -cost;
	float sintOverSini = eta;
	*wi = Vector3(sintOverSini * -wo.x, sintOverSini * -wo.y, cost);
	*pdf = 1.f;
	Spectrum F = fresnel.Evaluate(Frame::CosTheta(wo));

	// http://groups.google.com/group/pbrt/browse_thread/thread/6cf65c4132535685?pli=1

	// For radiance!!
	// return (et*et)/(ei*ei) * (Spectrum(1.f)-F) * T / fabsf(Frame::CosTheta(*wi));

	// For importance or power
	// When camera and light in the same media, it works well for radiance two...
	return /*(et*et)/(ei*ei) **/ (Spectrum(1.f)-F) * T / fabsf(Frame::CosTheta(*wi));
}

Spectrum OrenNayar::f(const Vector3 &wo, const Vector3 &wi) const {
	float sinthetai = Frame::SinTheta(wi);
	float sinthetao = Frame::SinTheta(wo);
	// cos(phi_i - phi_o)
	//= cos(phi_i)*cos(phi_o) + sin(phi_i)*sin(phi_o);
	float sinphii = Frame::SinPhi(wi), cosphii = Frame::CosPhi(wi);
	float sinphio = Frame::SinPhi(wo), cosphio = Frame::CosPhi(wo);
	float dcos = cosphii * cosphio + sinphii * sinphio;
	float maxcos = max(dcos, 0.f);
	// compute alpha and beta
	// alpha = max(theta_i, theta_o)
	// beta = min(theta_i, theta_o)
	float sinalpha, tanbeta;
	if(fabsf(Frame::CosTheta(wi)) > fabsf(Frame::CosTheta(wo))) {
		sinalpha = sinthetao;
		tanbeta = sinthetai / fabsf(Frame::CosTheta(wi));
	}
	else {
		sinalpha = sinthetai;
		tanbeta = sinthetao / fabsf(Frame::CosTheta(wo));
	}
	
	return R * INV_PI * (A + B * maxcos * sinalpha * tanbeta);
}

Spectrum Microfacet::f(const Vector3 &wo, const Vector3 &wi) const {
	float cosThetaO = Frame::AbsCosTheta(wo);
	float cosThetaI = Frame::AbsCosTheta(wi);
	if (cosThetaI == 0.f || cosThetaO == 0.f) return Spectrum(0.f);
	Vector3 wh = wi + wo;
	if (wh.x == 0. && wh.y == 0. && wh.z == 0.) return Spectrum(0.f);
	wh = Normalize(wh);
	float cosThetaH = Dot(wi, wh);
	Spectrum F = fresnel->Evaluate(cosThetaH);
	return R * distribution->D(wh) * G(wo, wi, wh) * F /
		(4.f * cosThetaI * cosThetaO);
}

Spectrum Microfacet::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &bxdfSample, float *pdf) const {
	distribution->Sample_f(wo, wi, bxdfSample, pdf);
	if(!Frame::SameHemisphere(wo, *wi))
		return Spectrum(0.f);
	return f(wo, *wi);
}

float Microfacet::Pdf(const Vector3 &wo, const Vector3 &wi) const {
	if(!Frame::SameHemisphere(wo, wi)) 
		return 0.f;
	return distribution->Pdf(wo, wi);
}

void Blinn::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &distributionSample, float *pdf) const {
	float costheta = powf(distributionSample.x, 1.f/(exponent+1));
	float sintheta = sqrtf(max(0.f, 1.f-costheta*costheta));
	float phi = distributionSample.y * 2.f * M_PI;
	Vector3 wh = SphericalDirection(sintheta, costheta, phi);
	if(!Frame::SameHemisphere(wo, wh)) 
		wh.z *= -1;
	*wi = -wo + 2.f * Dot(wo, wh) * wh;
	//float blinnPdf = ((exponent + 2.f) * powf(costheta, exponent)) / (2.f * M_PI * 4.f * Dot(wo, wh));
	float blinnPdf = ((exponent + 1.f) * powf(costheta, exponent)) / (2.f * M_PI * 4.f * Dot(wo, wh));
	if (Dot(wo, wh) <= 0.f) blinnPdf = 0.f;
	*pdf = blinnPdf;
}

float Blinn::Pdf(const Vector3 &wo, const Vector3 &wi) const {
	Vector3 wh = Normalize(wo + wi);
	float costheta = Frame::AbsCosTheta(wh);
	float blinnPdf = ((exponent + 1.f) * powf(costheta, exponent)) / (2.f * M_PI * 4.f * Dot(wo, wh));
	if (Dot(wo, wh) <= 0.f) blinnPdf = 0.f;
	return blinnPdf;
}

void Anisotropic::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &distributionSample, float *pdf) const {
	float phi, costheta;
	Point2 sample;
	float u1 = distributionSample.x, u2 = distributionSample.y;
	if (u1 < .25f) {
		sample.x = 4.f * u1;
		sample.y = u2;
		SampleFirstQuadrant(sample, &phi, &costheta);
	} else if (u1 < .5f) {
		sample.x = 4.f * (.5f - u1);
		sample.y = u2;
		SampleFirstQuadrant(sample, &phi, &costheta);
		phi = M_PI - phi;
	} else if (u1 < .75f) {
		sample.x = 4.f * (u1 - .5f);
		sample.y = u2;
		SampleFirstQuadrant(sample, &phi, &costheta);
		phi += M_PI;
	} else {
		sample.x = 4.f * (1.f - u1);
		sample.y = u2;
		SampleFirstQuadrant(sample, &phi, &costheta);
		phi = 2.f * M_PI - phi;
	}
	float sintheta = sqrtf(max(0.f, 1.f - costheta*costheta));
	Vector3 wh = SphericalDirection(sintheta, costheta, phi);
	if (!Frame::SameHemisphere(wo, wh)) wh = -wh;

	// Compute incident direction by reflecting about $\wh$
	*wi = -wo + 2.f * Dot(wo, wh) * wh;

	// Compute PDF for $\wi$ from anisotropic distribution
	float costhetah = Frame::AbsCosTheta(wh);
	float ds = 1.f - costhetah * costhetah;
	float anisotropic_pdf = 0.f;
	if (ds > 0.f && Dot(wo, wh) > 0.f) {
		float e = (ex * wh.x * wh.x + ey * wh.y * wh.y) / ds;
		float d = sqrtf((ex+1.f) * (ey+1.f)) * INV_TWOPI *
			powf(costhetah, e);
		anisotropic_pdf = d / (4.f * Dot(wo, wh));
	}
	*pdf = anisotropic_pdf;
}

float Anisotropic::Pdf(const Vector3 &wo, const Vector3 &wi) const {
	Vector3 wh = Normalize(wo + wi);
	// Compute PDF for $\wi$ from anisotropic distribution
	float costhetah = Frame::AbsCosTheta(wh);
	float ds = 1.f - costhetah * costhetah;
	float anisotropic_pdf = 0.f;
	if (ds > 0.f && Dot(wo, wh) > 0.f) {
		float e = (ex * wh.x * wh.x + ey * wh.y * wh.y) / ds;
		float d = sqrtf((ex+1.f) * (ey+1.f)) * INV_TWOPI *
			powf(costhetah, e);
		anisotropic_pdf = d / (4.f * Dot(wo, wh));
	}
	return anisotropic_pdf;
}

void Anisotropic::SampleFirstQuadrant(const Point2 &sample, float *phi, float *costheta) const {
	if (ex == ey)
		*phi = M_PI * sample.x * 0.5f;
	else
		*phi = atanf(sqrtf((ex+1.f) / (ey+1.f)) *
		tanf(M_PI * sample.x * 0.5f));
	float cosphi = cosf(*phi), sinphi = sinf(*phi);
	*costheta = powf(sample.y, 1.f/(ex * cosphi * cosphi +
		ey * sinphi * sinphi + 1));
}

Spectrum Lafortune::f(const Vector3 &wo, const Vector3 &wi) const {
	Spectrum ret = R * INV_PI;
	for(u_int i = 0; i < nLobes; i++) {
		Spectrum v = x[i] * wo.x * wi.x +
					y[i] * wo.y * wi.y +
					z[i] * wo.z * wi.z;
		ret += v.Pow(exponent[i]);
	}
	return ret;
}

Spectrum FresnelBlend::f(const Vector3 &wo, const Vector3 &wi) const {
	Spectrum diffuse = (28.f/(23.f*M_PI)) * Rd * (Spectrum(1.) - Rs) *
		(1 - powf(1 - .5f * fabsf(Frame::CosTheta(wi)),5)) *
		(1 - powf(1 - .5f * fabsf(Frame::CosTheta(wo)),5));
	Vector3 wh = Normalize(wi + wo);
	// what a fucking formula...
	Spectrum specular = distribution->D(wh) / 
						(8.f * M_PI * AbsDot(wi, wh) * 
						max(fabsf(Frame::CosTheta(wi)),fabsf(Frame::CosTheta(wo)))) *
						SchilckFresnel(Dot(wi, wh));
	return diffuse + specular;
}

Spectrum FresnelBlend::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &bxdfSample, float *pdf) const {
	float u1 = bxdfSample.x;
	float u2 = bxdfSample.y;
	Point2 sample;
	if(u1 < 0.5f) {
		sample.x = 2 * u1;
		sample.y = u2;
		*wi = CosineSamplingHemisphere(sample);
		if(wo.z < 0) wi->z *= -1;
	}
	else {
		sample.x = 2.f * (u1 - 0.5f);
		sample.y = u2;
		distribution->Sample_f(wo, wi, sample, pdf);
		if(!Frame::SameHemisphere(wo, *wi))
			return Spectrum(0.f);
	}
	*pdf  = Pdf(wo, *wi);
	return f(wo, *wi);
}

float FresnelBlend::Pdf(const Vector3 &wo, const Vector3 &wi) const {
	if(!Frame::SameHemisphere(wo, wi))
		return 0.f;
	return .5f * (fabsf(wi.z) * INV_PI + distribution->Pdf(wo, wi));
}

Spectrum Ward::f(const Vector3 &wo, const Vector3 &wi) const {
	
	if(!Frame::SameHemisphere(wi,wo))
		return Spectrum(0.f);

	Vector3 wh = Normalize(wo+wi);

	if(Dot(wh, wo) < 0.f)
		return Spectrum(0.f);
	
	Spectrum f = Rs * INV_FOURPI_ALPHAXY / sqrtf(wi.z * wo.z);
	float e = expf(-((wh.x*wh.x*invAlphaX2)+(wh.y*wh.y*invAlphaY2))/(wh.z*wh.z));

	return f*e;
}

Spectrum Ward::Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const {
	
	float phi = atanf(alphaY/alphaX*tanf(2*M_PI*BxDFSample.y));
	// make sure phi in the sample quadrant as 2 * M_PI * v
	float v =BxDFSample.y;
	if(v > 0.25 && v <= 0.5) // II
		phi += M_PI;
	else if(v > 0.5 && v < 0.75) // III
		phi -= M_PI;
	/*
	* In fact, we can compute as
	*	if( v > 0.25 && v <0.75)
	*		phi += M_PI
	* It makes no difference whether we add or minus M_PI 
	*/

	float cosphi = cosf(phi);
	float sinphi = sinf(phi);
	float thetah = atanf(sqrtf((-logf(BxDFSample.x))/((cosphi*cosphi*invAlphaX2)+(sinphi*sinphi*invAlphaY2))));
	
	Vector3 wh = SphericalDirection(sinf(thetah), cosf(thetah), phi);

	if(!Frame::SameHemisphere(wo, wh))
		wh.z *= -1;

	*wi = -wo + 2.f * Dot(wo, wh) * wh;

	*pdf = Pdf(wo, *wi);
	
	return f(wo, *wi);
}

float Ward::Pdf(const Vector3 &wo, const Vector3 &wi) const {

	if(!Frame::SameHemisphere(wo, wi))
		return (0.f);
	
	Vector3 wh = Normalize(wo+wi);

	if(Dot(wh, wo) < 0.f)
		return 0.f;

	float costhetah = Frame::CosTheta(wh);
	float f = INV_FOURPI_ALPHAXY / (costhetah*costhetah*costhetah) / Dot(wo,wh);
	float e = expf(-((wh.x*wh.x*invAlphaX2)+(wh.y*wh.y*invAlphaY2))/(wh.z * wh.z));

	return f*e;
}
